import pickle
import pandas as pd, numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer
import time
data = pickle.load(open('./train_tf_idf.pkl','rb'))
t1 = time.time()
train_word,train_classify,test_word,test_id,trn_term_doc,test_term_doc = data[0],data[1],data[2],data[3],data[4],data[5]
# list转成array形式
train_classify = np.array(train_classify)
# 字符变成int形式
y=(train_classify-1).astype(int)
clf = LogisticRegression(C=4, dual=True)
clf.fit(trn_term_doc, y)
preds=clf.predict_proba(test_term_doc)

#保存概率文件
test_prob=pd.DataFrame(preds)
test_prob.columns=["class_prob_%s"%i for i in range(1,preds.shape[1]+1)]
test_prob["id"]=test_id
test_prob.to_csv('.lr_probablity.csv',index=None)

#生成提交结果
preds=np.argmax(preds,axis=1)
test_pred=pd.DataFrame(preds)
test_pred.columns=["class"]
test_pred["class"]=(test_pred["class"]+1).astype(int)
print(test_pred.shape)
# print(test_id.shape)
test_pred["id"]=test_id
test_pred[["id","class"]].to_csv('./lr_result.csv',index=None)
t2=time.time()
print("time use:",t2-t1)